# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: wooyaggo

import json
import time
import random

import pyotp

from absl import app, flags

from selenium import webdriver
from selenium.webdriver.chrome.options import Options

FLAGS = flags.FLAGS

flags.DEFINE_string("cookie_file",
                    None,
                    "Path of output cookie_file which has authentication token.")

flags.DEFINE_string("userpass_key_file",
                    None,
                    "Path of json file which has owner/username/password fields.")

flags.DEFINE_boolean("go_coupon", None, "go coupon")

flags.DEFINE_string(
    "user_agent",
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/64.0.3282.119 Safari/537.36",
    "User Agent used for login.")

flags.DEFINE_string("proxy_server",
                    '',
                    "Proxy server. (e.g. socks5://127.0.0.1:12345 for ssh tunnel)")


def _create_driver(user_agent=None, proxy_server=None):
  opts = Options()
  if user_agent:
    opts.add_argument('user-agent=' + user_agent)

  if proxy_server:
    opts.add_argument('--proxy-server=' + proxy_server)

  driver = webdriver.Chrome('chromedriver', chrome_options=opts)

  driver.implicitly_wait(3)
  driver.get('https://bithumb.com')
  driver.implicitly_wait(8)
  driver.get('https://www.bithumb.com/u1/US101')
  driver.set_page_load_timeout(60)

  return driver


class BithumbChromeDriver(object):
  def __init__(self, user_agent=None, proxy_server=None):
    self.driver = _create_driver(user_agent, proxy_server)

  def login(self, userid, password, totp_secret_key):
    driver = self.driver

    driver.set_page_load_timeout(60)
    driver.implicitly_wait(10)
    # Input email and password.
    email_textbox = driver.find_element_by_xpath('//*[@id="user_id"]')
    password_textbox = driver.find_element_by_xpath('//*[@id="user_pw"]')
    email_textbox.send_keys(userid)
    password_textbox.send_keys(password)

    driver.implicitly_wait(3)
    time.sleep(random.randint(5, 8))
    driver.find_element_by_xpath('//*[@class="btn_org"]').click()

    # Input OTP and wait.
    time.sleep(random.randint(5, 8))
    otp = pyotp.TOTP(totp_secret_key).now()
    otp_textbox = driver.find_element_by_xpath('//*[@id="otpCertify"]')
    otp_textbox.send_keys(otp)
    driver.find_element_by_xpath('//*[@id="otpOK"]').click()
    time.sleep(5)

  def go_coupon(self):
    driver = self.driver

    driver.get('https://www.bithumb.com/u5/US505')
    time.sleep(30)

  def get_cookies(self):
    return self.driver.get_cookies()

  def quit(self):
    if self.driver is not None:
      self.driver.quit()
      self.driver = None


def main(_):
  with open(FLAGS.userpass_key_file) as f:
    userpass = json.load(f)

  assert 'owner' in userpass
  assert 'username' in userpass
  assert 'password' in userpass
  assert 'totp_secret_key' in userpass

  try:
    driver = BithumbChromeDriver(FLAGS.user_agent, FLAGS.proxy_server)
    driver.login(userpass['username'], userpass['password'], userpass['totp_secret_key'])
    cookie = driver.get_cookies()

    if FLAGS.go_coupon:
      driver.go_coupon()
  except Exception:
    raise
  finally:
    driver.quit()

  res = json.dumps(cookie, sort_keys=True, indent=4, separators=(',', ': '))
  print(res)
  if FLAGS.cookie_file is not None:
    erroneous_cookie_file = []
    for cookie_file in FLAGS.cookie_file.split(','):
      try:
        print(cookie_file)
        with open(cookie_file, 'w') as f:
          f.write(res)
      except Exception:
        print('failed to save to', cookie_file)
        erroneous_cookie_file.append(cookie_file)

    if len(erroneous_cookie_file):
      raise Exception('Failed to save to: ' + repr(erroneous_cookie_file))


if __name__ == '__main__':
  app.run(main)
